/**
 * 
 */
package edu.umd.clip.lm.model;

import java.io.*;
import java.util.*;

import com.sleepycat.bind.tuple.TupleInput;
import com.sleepycat.bind.tuple.TupleOutput;

import edu.umd.clip.lm.util.*;

/**
 * @author Denis Filimonov <den@cs.umd.edu>
 *
 */
public class OnDiskCompactProbTree implements Cloneable {
	private int[] compactHiddenFactors;
	private float[] probabilities;
	private short size;
	private double scale;
	
	public OnDiskCompactProbTree(int compactFactors, double prob) {
		this(1);
		addItem(compactFactors, (float) prob);
	}
	
	public OnDiskCompactProbTree(int size) {
		compactHiddenFactors = new int[size];
		probabilities = new float[size];
		this.size = 0;
		this.scale = 1;
	}

	/**
	 * @param compactHiddenFactors
	 * @param probabilities
	 * @param scale
	 */
	public OnDiskCompactProbTree(int[] compactHiddenFactors, float[] probabilities, double scale) {
		assert(probabilities.length == compactHiddenFactors.length);
		
		this.compactHiddenFactors = compactHiddenFactors;
		this.probabilities = probabilities;
		this.scale = scale;
		this.size = (short) probabilities.length;
	}

	private void addItem(int compactFactors, float prob) {
		int pos = Arrays.binarySearch(compactHiddenFactors, 0, size, compactFactors);
		if (pos >= 0) {
			probabilities[pos] += prob;
		} else {
			pos = -pos - 1;
			if (size <= compactHiddenFactors.length) {
				grow(size+1);
			}
			if (pos < size) {
				System.arraycopy(compactHiddenFactors, pos, compactHiddenFactors, pos+1, size-pos);
				System.arraycopy(probabilities, pos, probabilities, pos+1, size-pos);
			}
			compactHiddenFactors[pos] = compactFactors;
			probabilities[pos] = prob;
			++size;
		}
	}
	
	private void grow(int desiredSize) {
		if (desiredSize > compactHiddenFactors.length) {
			int newSize = compactHiddenFactors.length << 1;
			while(newSize <= desiredSize) {
				newSize <<= 1;
			}
			compactHiddenFactors = Arrays.copyOf(compactHiddenFactors, newSize);
			probabilities = Arrays.copyOf(probabilities, newSize); 
		}
	}
	
	public OnDiskCompactProbTree(ObjectInput in) throws IOException, ClassNotFoundException {
		readExternal(in);
	}
	
	/**
	 * @param in
	 */
	public OnDiskCompactProbTree(TupleInput in) throws IndexOutOfBoundsException {
		readTuple(in);
	}

	public static OnDiskCompactProbTree merge(OnDiskCompactProbTree t1, OnDiskCompactProbTree t2) {
		long start = Timing.nanoTime();
		try {
			if (t1.getSize() < t2.getSize()){
				OnDiskCompactProbTree tmp = t1;
				t1 = t2;
				t2 = tmp;
			}
			
			if (t2.scale == 0) {
				return (OnDiskCompactProbTree) t1.clone();
			} else if (t1.scale == 0) {
				return (OnDiskCompactProbTree) t2.clone();
			}
			// at this point, t1 is the bigger tree
			// normally this also means that it includes all elements from the smaller tree
			// as t1 would represent the distribution at the parent node in the decision tree 
			OnDiskCompactProbTree result = new OnDiskCompactProbTree(t1.getSize());
			final int size1 = t1.getSize();
			final int size2 = t2.getSize();
			final int factors1[] = t1.compactHiddenFactors;
			final int factors2[] = t2.compactHiddenFactors;
			final double scale2 = (t2.scale / t1.scale);
			int point1 = 0;
			int point2 = 0;
			int out = 0;
			while(point1 < size1 || point2 < size2) {
				if (out == result.compactHiddenFactors.length) {
					result.grow(out + Math.max(size1-point1, size2-point2));
				}
				if (point2 == size2) {
					result.compactHiddenFactors[out] = factors1[point1];
					result.probabilities[out] = t1.probabilities[point1];
					point1++;
				} else if (point1 < size1) {
					if (factors1[point1] < factors2[point2]) {
						result.compactHiddenFactors[out] = factors1[point1];
						result.probabilities[out] = t1.probabilities[point1];
						point1++;					
					} else if (factors1[point1] == factors2[point2]) {
						result.compactHiddenFactors[out] = factors2[point2];
						result.probabilities[out] = (float) (t1.probabilities[point1] + t2.probabilities[point2] * scale2);
						point1++;
						point2++;										
					} else {
						result.compactHiddenFactors[out] = factors2[point2];
						result.probabilities[out] = (float) (t2.probabilities[point2] * scale2);
						point2++;															
					}
				} else {
					// point1 == size1 && point2 < size2
					result.compactHiddenFactors[out] = factors2[point2];
					result.probabilities[out] = (float) (t2.probabilities[point2] * scale2);
					point2++;
				}
				out++;
			}
			result.size = (short) out;
			result.scale = t1.scale;
			return result;
		} finally {
			long end = Timing.nanoTime();
			ProbTree.totalLeafyMergeCount += 1;
			ProbTree.totalLeafyMergeTime += end - start;
		}
	}
	
	/* (non-Javadoc)
	 * @see edu.umd.clip.lm.model.CompactProbTreeInterface#checkTree()
	 */
	public boolean checkTree() {
		return true;
	}

	/* (non-Javadoc)
	 * @see edu.umd.clip.lm.model.CompactProbTreeInterface#compact()
	 */
	public void compact() {
		throw new UnsupportedOperationException();
	}

	public double getProbability(int compactFactors) {
		int pos = Arrays.binarySearch(compactHiddenFactors, compactFactors);
		if (pos < 0 || pos >= size) return 0;
		return probabilities[pos] * scale;
	}

	/* (non-Javadoc)
	 * @see edu.umd.clip.lm.model.CompactProbTreeInterface#getTotalProb()
	 */
	public double getTotalProb() {
		double prob = 0;
		for(int i=0; i<size; ++i) {
			prob += probabilities[i];
		}
		return prob * scale;
	}

	/* (non-Javadoc)
	 * @see edu.umd.clip.lm.model.CompactProbTreeInterface#prune(double)
	 */
	public AbstractCompactProbTree prune(double threshold) {
		throw new UnsupportedOperationException();
		//return this;
	}

	/* (non-Javadoc)
	 * @see edu.umd.clip.lm.model.CompactProbTreeInterface#readExternal(java.io.ObjectInput)
	 */
	public void readExternal(ObjectInput in) throws IOException, ClassNotFoundException {
		size = in.readShort();
		
		compactHiddenFactors = new int[size];
		probabilities = new float[size];
		for(int i=0; i<size; ++i) {
			compactHiddenFactors[i] = in.readInt();
			probabilities[i] = in.readFloat();
		}
		scale = in.readDouble();
	}
	
	public void readTuple(TupleInput in) throws IndexOutOfBoundsException {
		size = in.readShort();
		
		compactHiddenFactors = new int[size];
		probabilities = new float[size];
		for(int i=0; i<size; ++i) {
			compactHiddenFactors[i] = in.readInt();
			probabilities[i] = in.readFloat();
		}
		scale = in.readDouble();
	}

	/* (non-Javadoc)
	 * @see edu.umd.clip.lm.model.CompactProbTreeInterface#scale(double)
	 */
	public void scale(double scale) {
		this.scale *= scale;
	}

	public void normalize() {
		double treeProb = 0;
		for(int i=0; i<size; ++i) {
			treeProb += probabilities[i];
		}
		if (treeProb > 0) {
			for(int i=0; i<size; ++i) {
				probabilities[i] /= treeProb;
			}
			scale *= treeProb;
		}
	}
	
	public void setProbability(int compactFactors, double prob) {
		prob /= scale;
		int pos = Arrays.binarySearch(compactHiddenFactors, compactFactors);
		if (pos >=0 && pos < size) {
			probabilities[pos] = (float) prob;
		} else {
			addItem(compactFactors, (float) prob);
		}
	}

	public void writeExternal(ObjectOutput out) throws IOException {
		out.writeShort(size);
		for(int i=0; i<size; ++i) {
			out.writeInt(compactHiddenFactors[i]);
			out.writeFloat(probabilities[i]);
		}
		out.writeDouble(scale);
	}

	public void writeTuple(TupleOutput out) {
		out.writeShort(size);
		for(int i=0; i<size; ++i) {
			out.writeInt(compactHiddenFactors[i]);
			out.writeFloat(probabilities[i]);
		}
		out.writeDouble(scale);
	}

	@Override
	public Object clone() {
		long start = Timing.nanoTime();
		try {
			OnDiskCompactProbTree newTree = new OnDiskCompactProbTree(size);
			System.arraycopy(compactHiddenFactors, 0, newTree.compactHiddenFactors, 0, size);
			System.arraycopy(probabilities, 0, newTree.probabilities, 0, size);
			newTree.scale = scale;
			newTree.size = size;
			return newTree;
		} finally {
			long end = Timing.nanoTime();
			ProbTree.totalCloningTime += end - start;
		}
	}

	public double getProb() {
		return getTotalProb();
	}

	public void addProbability(int packedHiddenFactors, double prob) {
		addItem(packedHiddenFactors, (float) (prob/scale));
		
	}

	public int[] getCompactHiddenFactors() {
		return compactHiddenFactors;
	}

	public float[] getProbabilities() {
		return probabilities;
	}

	public short getSize() {
		return size;
	}

	public double getScale() {
		return scale;
	}

	private static class SizeComparator implements Comparator<OnDiskCompactProbTree> {
		/* (non-Javadoc)
		 * @see java.util.Comparator#compare(java.lang.Object, java.lang.Object)
		 */
		@Override
		public int compare(OnDiskCompactProbTree o1, OnDiskCompactProbTree o2) {
			// inverse order
			return o2.size - o1.size;
		}
	}
	private static final SizeComparator sizeComparator = new SizeComparator();
	/**
	 * @param probTrees
	 * @return
	 */
	public static OnDiskCompactProbTree merge(ArrayList<OnDiskCompactProbTree> probTrees) {
		if (probTrees.size() == 2) {
			return merge(probTrees.get(0), probTrees.get(1));
		}
		// sort the trees by their size, the largest first
		OnDiskCompactProbTree trees[] = probTrees.toArray(new OnDiskCompactProbTree[probTrees.size()]);
		Arrays.sort(trees, sizeComparator);
		OnDiskCompactProbTree tree = merge(trees[0], trees[1]);
		for(int i=2; i<trees.length; ++i) {
			tree = merge(tree, trees[i]);
		}
		return tree;
	}

	public static OnDiskCompactProbTree merge(OnDiskCompactProbTree probTrees[], double weights[]) {
		if (probTrees.length < 2) {
			if (probTrees.length == 0) {
				return null;
			}
			OnDiskCompactProbTree result = (OnDiskCompactProbTree) probTrees[0].clone();
			result.scale(weights[0]);
			return result;
		}
		
		int largestTreePos = 0;
		short largestTreeSize = 0;
		for(int i=0; i<probTrees.length; ++i) {
			short mySize = probTrees[i].size; 
			if (mySize > largestTreeSize) {
				largestTreeSize = mySize;
				largestTreePos = i;
			}
		}
		
		CompactReadOnlyInt2FloatHashMap probs = new CompactReadOnlyInt2FloatHashMap(probTrees[largestTreePos].compactHiddenFactors, probTrees[largestTreePos].probabilities);
		
		float values[] = probs.values();
		
		// scale the values
		{
			float scaleFactor = (float) (weights[largestTreePos] * probTrees[largestTreePos].scale);
			for(int i=0; i<values.length; ++i) {
				values[i] *= scaleFactor;
			}
		}
		
		// add the other trees
		CompactReadOnlyIntHashSet keys = probs.getKeysHashSet();
		for(int tree=0; tree<probTrees.length; ++tree) {
			if (tree != largestTreePos) {
				OnDiskCompactProbTree myTree = probTrees[tree];
				float scalingFactor = (float) (myTree.getScale() * weights[tree]); 
				for(int i=0; i<myTree.size; ++i) {
					int idx = keys.getKeyPosition(myTree.compactHiddenFactors[i]);
					assert(idx >= 0);
					// we assume that the largest tree contains ALL possible tags for the word
					values[idx] += myTree.probabilities[i] * scalingFactor;
				}
			}
		}
		
		return new OnDiskCompactProbTree(probs.keys(), probs.values(), 1.0);
	}
}
